package pers.chao.tool.thread.pool.impl;

import org.slf4j.Logger;
import pers.chao.tool.infrastructure.log.BaseLogFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.*;

/**
 * Description: 线程池队列配置和拒绝策略异常配置
 *
 * @author W.Chao
 * @date 2020/12/13 22:45
 **/
public final class ExceedWait {
	public static class RejectedHandler implements RejectedExecutionHandler {
		public static final Logger logger = BaseLogFactory.getLogger(RejectedHandler.class);

		@Override
		public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
			if (!executor.isShutdown()) {
				logger.warn("[ExceedWaitRejectedHandler] put the rejected command to the exceed queue" +
						" in the rejectedExecution method: {}", r);
				((Queue) executor.getQueue()).putExceed(r, executor);
			}
		}
	}

	/**
	 * 自定义实现的任务队列
	 */
	public static class Queue extends SynchronousQueue<Runnable> {
		public static final Logger logger = BaseLogFactory.getLogger(Queue.class);
		private static final long serialVersionUID = -7249478138264167124L;
		private final LinkedBlockingQueue<Runnable> exceedQueue = new LinkedBlockingQueue<>();

		public Queue() {
			super(true);
		}

		@Override
		public boolean offer(Runnable runnable) {
			logger.info("[ExceedWaitQueue] offer() called with: runnable = {}", runnable);
			boolean result = super.offer(runnable);
			logger.info("[ExceedWaitQueue] offer() returned: {}", result);
			return result;
		}

		@Override
		public Runnable poll(long timeout, TimeUnit unit) throws InterruptedException {
			logger.info("[ExceedWaitQueue] poll() called with: timeout = {}, unit = {}", timeout, unit);
			// 查询主队列是否有堆积的任务
			Runnable result = super.poll();
			// 主队列没有任务且等待队列有任务时从等待队列(辅助)中获取任务
			if (exceedQueue.size() > 0 && result == null) {
				result = exceedQueue.poll();
			}
			// 两个队列都没有任务时，等待其他线程插入任务直到超时
			if (result == null) {
				result = super.poll(timeout, unit);
			}

			logger.info("[ExceedWaitQueue] poll() returned: {}", result);
			return result;
		}

		@Override
		public Runnable take() throws InterruptedException {
			logger.info("[ExceedWaitQueue] take() called");
			// 检查主队列是否有任务
			Runnable result = super.poll();
			// 主队列没有任务且等待队列有任务时从等待队列(辅助)中获取任务
			if (result == null && exceedQueue.size() > 0) {
				result = exceedQueue.poll();
			}

			// 两个队列都没有任务时一直等待至有任务
			if (result == null) {
				result = super.take();
			}

			logger.info("[ExceedWaitQueue] take() returned: {}", result);
			return result;
		}

		@Override
		public boolean remove(Object o) {
			logger.info("[ExceedWaitQueue] remove() called with: o = [{}]", o);
			return exceedQueue.remove(o);
		}

		@Override
		public boolean removeAll(Collection<?> c) {
			logger.info("[ExceedWaitQueue] removeAll() called with: c = [{}]", c);
			return exceedQueue.removeAll(c);
		}

		public int exceedSize() {
			return exceedQueue.size();
		}

		public void putExceed(Runnable e, ThreadPoolExecutor executor) {
			exceedQueue.offer(e);

			final int activeCount = executor.getActiveCount();
			if (activeCount <= 0) {
				// 线程池没有保持固定数量的可活动核心线程，且主队列等待线程插入任务
				// 需要将等待辅队列的任务唤醒线程池
				final Runnable next = exceedQueue.poll();
				if (next != null) {
					logger.info("[ExceedWaitQueue] putExceed and activeCount({}), need to " +
							"wake up the pool with next({})", activeCount, next);
					executor.execute(next);
				}
			}
		}

		/**
		 * 将队列中的任务全部清空，并返回所有队列中的任务
		 *
		 * @return
		 */
		public List<Runnable> drainExceedQueue() {
			BlockingQueue<Runnable> q = exceedQueue;
			ArrayList<Runnable> taskList = new ArrayList<>();
			q.drainTo(taskList);
			if (!q.isEmpty()) {
				for (Runnable r : q.toArray(new Runnable[0])) {
					if (q.remove(r))
						taskList.add(r);
				}
			}
			return taskList;
		}
	}
}
